from paddlenlp.transformers import ErnieForMaskedLM, ErnieTokenizer
import paddle
import random

class MLMAugmenter:
    def __init__(self, model_name="ernie-1.0", mask_ratio=0.15):
        self.tokenizer = ErnieTokenizer.from_pretrained(model_name)
        self.model = ErnieForMaskedLM.from_pretrained(model_name)
        self.mask_ratio = mask_ratio
    
    def augment(self, text, n=1):
        """生成n个增强样本"""
        augmented_texts = []
        for _ in range(n):
            augmented_texts.append(self._augment_once(text))
        return augmented_texts if n > 1 else augmented_texts[0]
    
    def _augment_once(self, text):
        # 分词
        tokens = self.tokenizer.tokenize(text)
        if len(tokens) <= 2:  # 至少保留[CLS]和[SEP]
            return text
        
        # 计算掩码数量
        n_mask = max(1, int(len(tokens) * self.mask_ratio))
        n_mask = min(n_mask, len(tokens) - 2)  # 确保至少保留[CLS]和[SEP]
        
        # 随机选择位置进行掩码
        mask_indices = random.sample(range(1, len(tokens)-1), n_mask)
        
        # 创建掩码输入
        masked_tokens = tokens.copy()
        for idx in mask_indices:
            masked_tokens[idx] = '[MASK]'
        
        # 转换为模型输入
        inputs = self.tokenizer.convert_tokens_to_ids(masked_tokens)
        inputs = paddle.to_tensor([inputs])
        
        # 预测掩码位置的词
        with paddle.no_grad():
            outputs = self.model(inputs)
        
        # 获取预测结果
        predictions = outputs[0]
        
        # 替换掩码位置的词
        augmented_tokens = tokens.copy()
        for idx in mask_indices:
            # 可选：使用top-k采样而非argmax，增加多样性
            # pred_idx = paddle.argsort(predictions[0, idx], descending=True)[:5]
            # pred_idx = random.choice(pred_idx).item()
            pred_idx = paddle.argmax(predictions[0, idx]).item()
            augmented_tokens[idx] = self.tokenizer.convert_ids_to_tokens([pred_idx])[0]
        
        # 转换回文本
        augmented_text = self.tokenizer.convert_tokens_to_string(augmented_tokens)
        return augmented_text

# 使用示例
augmenter = MLMAugmenter()
text = "这是一个关于自然语言处理的示例句子。"
augmented_text = augmenter.augment(text)
print(f"原文: {text}")
print(f"增强后: {augmented_text}")